import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix
import numpy as np
import matplotlib
import tensorflow as tf
from tqdm import tqdm
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.models import load_model
import json
import time
import os

font = {'family': 'SimHei', "size": 24}
matplotlib.rc('font', **font)
input_size = 128


def get_confusion_matrix(label, y_pred, cls):
    Confusion_Matrix = confusion_matrix(label, y_pred)
    cls_sum = Confusion_Matrix.sum(axis=0)
    Confusion_Matrix = Confusion_Matrix / cls_sum  # 归一化
    print(np.around(Confusion_Matrix, 3))
    Confusion_Matrix = Confusion_Matrix * 255
    plt.matshow(Confusion_Matrix, cmap=plt.cm.Greens)
    plt.xlabel('predict')
    plt.ylabel('true')
    xlocations = np.array(range(len(cls)))
    print(cls)
    print(xlocations)
    plt.xticks(xlocations, cls, rotation=0)
    plt.yticks(xlocations, cls)
    plt.show()


def tflite_pre(modelpath, dataset_root, batch_size=128, input_size=96):
    test_root = dataset_root

    valid_generator = ImageDataGenerator()
    valid_dataset = valid_generator.flow_from_directory(batch_size=batch_size,
                                                        directory=test_root,
                                                        shuffle=True,
                                                        target_size=(input_size, input_size))
    interpreter = tf.lite.Interpreter(model_path=modelpath)

    input_index = interpreter.get_input_details()
    output_index = interpreter.get_output_details()

    acc = 0
    sum = 0
    labels = np.array([])
    predictions = np.array([])
    with tqdm(valid_dataset, desc='test_tflite') as d:
        for batch_number, (img, label) in enumerate(d):
            batch = len(label)
            label = label.argmax(axis=1)
            sum += batch
            interpreter.resize_tensor_input(input_index[0]['index'], (batch, input_size, input_size, 3))
            interpreter.resize_tensor_input(output_index[0]['index'], (batch, 15))
            interpreter.allocate_tensors()
            interpreter.set_tensor(input_index[0]['index'], np.array(img, dtype=np.float32))
            interpreter.invoke()
            output = interpreter.get_tensor(output_index[0]['index'])
            output = output.argmax(axis=1)
            correct_pre = np.where(output == label, 1, 0)
            acc += correct_pre.sum()

            predictions = np.append(predictions, output)
            labels = np.append(labels, label)

            if sum >= len(valid_dataset) * batch_size:
                break
            d.set_postfix(acc=acc / sum)
            d.update()
    print("小类准确率：", acc / sum)
    get_confusion_matrix(labels, predictions, valid_dataset.class_indices)
    return acc / sum


def convert_to_tf_lite(model_path, valid_input_size, valid_path=None):
    model = load_model(model_path)

    converter = tf.lite.TFLiteConverter.from_keras_model(model)
    converter.post_training_quantize = True
    tflite_model = converter.convert()

    start_time = time.strftime('%m_%d_%H_%M')
    name = model_path.split('/')[-1][:-3]

    save_root = "./tflite_model/%s_" % start_time + '_' + name
    os.mkdir(save_root)

    save_path = save_root + "/" + str(valid_input_size) + name + '.tflite'

    with open(save_path, 'wb') as f:
        f.write(tflite_model)
    if valid_path is not None:
        acc = tflite_pre(modelpath=save_path,
                         dataset_root=valid_path,
                         input_size=valid_input_size)
        remark = {'model_path': model_path, 'input_size': valid_input_size, 'valid_path': valid_path, 'acc': acc}
        with open(save_root + "/remark.json", 'w', encoding='utf-8') as f:
            json.dump(remark, f)


def batch_test_models(model_paths, dataset_root, input_size=128, batch_size=128):
    valid_generator = ImageDataGenerator()
    valid_dataset = valid_generator.flow_from_directory(batch_size=batch_size,
                                                        directory=dataset_root,
                                                        shuffle=False,
                                                        target_size=(input_size, input_size))

    class_indices = valid_dataset.class_indices
    inv_class_indices = {v: k for k, v in class_indices.items()}

    for model_path in model_paths:
        interpreter = tf.lite.Interpreter(model_path=model_path)
        input_index = interpreter.get_input_details()
        output_index = interpreter.get_output_details()

        print(f"Testing model: {model_path}")
        num_samples = len(valid_dataset.filenames)
        num_batches = int(np.ceil(num_samples / batch_size))

        for batch_idx in range(num_batches):
            start_idx = batch_idx * batch_size
            end_idx = min(start_idx + batch_size, num_samples)
            batch_images = []
            batch_labels = []

            for i in range(start_idx, end_idx):
                img_path = os.path.join(dataset_root, valid_dataset.filenames[i])
                img = tf.keras.preprocessing.image.load_img(img_path, target_size=(input_size, input_size))
                img = tf.keras.preprocessing.image.img_to_array(img)
                img = np.expand_dims(img, axis=0)
                batch_images.append(img)
                true_label = valid_dataset.classes[i]
                batch_labels.append(true_label)

            batch_images = np.vstack(batch_images)
            batch_labels = np.array(batch_labels)

            interpreter.resize_tensor_input(input_index[0]['index'], (len(batch_images), input_size, input_size, 3))
            interpreter.resize_tensor_input(output_index[0]['index'], (len(batch_images), 15))
            interpreter.allocate_tensors()
            interpreter.set_tensor(input_index[0]['index'], np.array(batch_images, dtype=np.float32))
            interpreter.invoke()
            output = interpreter.get_tensor(output_index[0]['index'])
            predictions = np.argmax(output, axis=1)

            num_images = len(batch_images)
            images_per_row = 4
            images_per_col = 2
            num_display = images_per_row * images_per_col
            num_display = min(num_display, num_images)

            fig, axes = plt.subplots(images_per_col, images_per_row, figsize=(images_per_row * 3, images_per_col * 3))
            axes = axes.flatten()

            for i in range(num_display):
                true_class = inv_class_indices[batch_labels[i]]
                pred_class = inv_class_indices[predictions[i]]
                is_misclassified = batch_labels[i] != predictions[i]

                axes[i].imshow(batch_images[i].astype(np.uint8))
                if is_misclassified:
                    # 添加红色文字标注
                    axes[i].text(0.5, 0.5, "预测错误", transform=axes[i].transAxes, 
                                color='red', fontsize=18, ha='center', va='center')

                title_text = f"True: {true_class}\nPred: {pred_class}"
                axes[i].set_title(title_text, fontsize=12)
                axes[i].axis('off')

            for j in range(num_display, len(axes)):
                axes[j].axis('off')

            plt.tight_layout(pad=2)
            plt.show()


# 示例使用（请替换为实际路径）
model_paths = ['./test.tflite']
dataset_root = './0-99' 
batch_test_models(model_paths, dataset_root, input_size=input_size)